package com.ruoyi.framework.security.handle;

import com.ruoyi.framework.config.OAuth2CacheManager;
import jakarta.annotation.Resource;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.security.core.Authentication;
import org.springframework.security.oauth2.client.authentication.OAuth2AuthenticationToken;
import org.springframework.security.oauth2.core.oidc.user.OidcUser;
import org.springframework.security.oauth2.server.resource.authentication.JwtAuthenticationToken;
import org.springframework.security.web.authentication.logout.SecurityContextLogoutHandler;
import org.springframework.stereotype.Component;

import java.util.Map;

@Component
@Slf4j
public class OAuth2LogoutHandler extends SecurityContextLogoutHandler {

    @Resource
    private OAuth2CacheManager oauth2CacheManager;

    @Override
    public void logout(HttpServletRequest request, HttpServletResponse response, Authentication authentication) {
        String sid = null;
        if (authentication instanceof OAuth2AuthenticationToken) {
            OidcUser oidcUser = (OidcUser) authentication.getPrincipal();
            sid = oidcUser.getAttribute("sid");
        } else if (authentication instanceof JwtAuthenticationToken jwtToken) {
            Map<String, Object> claims = jwtToken.getToken().getClaims();
            sid = claims.containsKey("sid") ? claims.get("sid").toString() : null;
        }
        String idToken = oauth2CacheManager.getIdToken(request, sid);
        request.setAttribute("idToken", idToken);
        request.setAttribute("sessionId", request.getSession().getId());
        oauth2CacheManager.clearCache(request, sid);
        super.logout(request, response, authentication);
    }
}